import os
import random
from typing import List

import torch
from PIL import Image

from data.base_dataset import BaseDataset
import abc


# Construct supervised image pairs in the form of pix2pix dataset
class Pix2pixDataset(BaseDataset):
    def __init__(self, opt):
        super(Pix2pixDataset, self).__init__(opt)
        labels = self.get_label_pt()
        self.labels = labels['labels']
        self.groups = labels['groups']
        self.idxs = list(range(len(self.labels)))
        sketch_paths, image_paths = self.get_paths()

        if not opt.no_pairing_check:
            for path1, path2 in zip(sketch_paths, image_paths):
                assert self.paths_match(path1, path2), \
                    "The label-image pair (%s, %s) do not look like the right pair because the filenames are quite different. Are you sure about the pairing? Please see data/pix2pix_dataset.py to see what is going on, and use --no_pairing_check to bypass this." % (
                        path1, path2)

        self.sketch_paths = sketch_paths
        self.image_paths = image_paths
        size = len(self.idxs)
        self.dataset_size = size
        # self.hard_reference_probability = 0 if opt.phase == 'test' else opt.hard_reference_probability
        self.ref_dict = self.get_ref(self.image_paths, self.labels, self.groups)

    @abc.abstractmethod
    def get_paths(self):
        pass

    @abc.abstractmethod
    def get_label_pt(self):
        pass

    def paths_match(self, path1, path2):
        filename1_without_ext = os.path.splitext(os.path.basename(path1))[0]
        filename2_without_ext = os.path.splitext(os.path.basename(path2))[0]
        return filename1_without_ext == filename2_without_ext

    def get_sketch_tensor(self, path):
        label = Image.open(path).convert('L')
        params1 = self.get_params(label.size)
        transform_label = self.get_sketch_transform(params1, method=Image.NEAREST, normalize=True)
        label_tensor = transform_label(label)
        # label_tensor = transform_label(label) * 255.0
        # label_tensor[label_tensor == 255] = self.opt.label_nc  # 'unknown' is opt.label_nc
        return label_tensor, params1, transform_label

    def get_ref(self, image_paths: List, labels: List, groups: List) -> dict:
        ref_dict = {}
        idxs = self.idxs
        for idx, img in enumerate(image_paths):
            label = labels[idx]
            group = groups[label]
            similar_color_img_idx = random.choice(group)
            hard_img_idx = random.choice(idxs)
            val = [similar_color_img_idx, hard_img_idx]
            ref_dict[img] = val
        return ref_dict

    @abc.abstractmethod
    def postprocess(self, input_dict):
        pass

    @abc.abstractmethod
    def hardsample(self):
        pass

    # -------------------------torch dataset implementation-------------------------#
    def __len__(self):
        return self.dataset_size

    def __getitem__(self, idx1):
        ref = self.image_paths[idx1]
        ref_label = self.sketch_paths[idx1]
        similar_color_img_idx, hard_img_idx = self.ref_dict[ref]
        similar_color_img = self.image_paths[similar_color_img_idx]
        similar_color_img_label = self.sketch_paths[similar_color_img_idx]
        hard_img = self.image_paths[hard_img_idx]
        hard_img_label = self.sketch_paths[hard_img_idx]

        ref = Image.open(ref).convert("RGB")
        similar_color_img = Image.open(similar_color_img).convert("RGB")
        hard_img = Image.open(hard_img).convert("RGB")

        ref_label, params1, label_transform = self.get_sketch_tensor(ref_label)
        similar_color_img_label = Image.open(similar_color_img_label).convert('L')
        hard_img_label = Image.open(hard_img_label).convert('L')
        similar_color_img_label = label_transform(similar_color_img_label)
        hard_img_label = label_transform(hard_img_label)

        transform_image = self.get_transform(params1, normalize=True)
        ref = transform_image(ref)
        similar_color_img = transform_image(similar_color_img)
        hard_img = transform_image(hard_img)

        self_ref_flag = torch.ones_like(ref)

        input_dict={
            'label': hard_img_label,
            'image': hard_img,
            'self_ref': self_ref_flag,
            'ref': ref,
            'label_ref': ref_label,
            'similar_ref': similar_color_img
        }

        input_dict = self.postprocess(input_dict)
        return input_dict




    # def __getitem__(self, index):
    #     # semantic label
    #     sketch_path = self.sketch_paths[index]
    #     sketch_tensor, params1, sketch_transform = self.get_sketch_tensor(sketch_path)
    #     # input image (real images)
    #     image_path = self.image_paths[index]
    #
    #     if not self.opt.no_pairing_check:
    #         assert self.paths_match(sketch_path, image_path), \
    #             "The label_path %s and image_path %s don't match." % \
    #             (sketch_path, image_path)
    #
    #     image = Image.open(image_path).convert("RGB")
    #     transform_image = self.get_transform(params1, normalize=True)
    #     image_tensor = transform_image(image)
    #
    #     if self.hardsample() or self.opt.phase == 'test':
    #         ref_path = self.ref_dict[image_path][1]
    #     else:
    #         ref_path = self.ref_dict[image_path][0]
    #
    #     ref_img = Image.open(ref_path).convert('RGB')
    #     ref_tensor = transform_image(ref_img)
    #     self_ref_flag = torch.ones_like(ref_tensor)
    #
    #     sketch_ref = [sketch_path for sketch_path in self.sketch_paths if os.path.basename(sketch_path) == os.path.basename(ref_path)]
    #     assert len(sketch_ref) is not 0
    #     sketch_ref = sketch_ref[0]
    #     sketch_ref = Image.open(sketch_ref).convert('L')
    #     sketch_ref_tensor = sketch_transform(sketch_ref)
    #
    #     input_dict = {'label': sketch_tensor,
    #                   'image': image_tensor,
    #                   'path': image_path,
    #                   'self_ref': self_ref_flag,
    #                   'ref': ref_tensor,
    #                   'label_ref': sketch_ref_tensor
    #                   }
    #     input_dict = self.postprocess(input_dict)
    #     return input_dict
